# Dickstein, Ho, and Mark (2023)
# This script adjusts the data and creates the necessary 
# variable names, depending on the type of estimation. 

if(!est_type %in% c("full_only", "switcher_only", "mixed")){
  stop("Invalid est_type")
}

if(est_type == "full_only"){
  # Extract the variable names:
  W1_names <- var_names$W1_names
  W2_names <- var_names$W2_names
  W3_names <- var_names$W3_names
  Z_names <- var_names$Z_names
  
  # Data is left as is. 
}

if(est_type == "switcher_only"){
  # Extract variable names:
  W1_names <- var_names$W1_names
  W2_names <- var_names$W2_names
  W3_names <- var_names$W3_names
  Z_names <- var_names$Z_names
  
  # Data is subsetted to switchers: 
  data <- data[switcher == 1]
}

if(est_type == "mixed"){
  # Create new variables name0, name1 (depending if switcher): 
  W1_switch <- var_names$W1_names[var_switch$W1_names == 1]
  W2_switch <- var_names$W2_names[var_switch$W2_names == 1]
  W3_switch <- var_names$W3_names[var_switch$W3_names == 1]
  Z_switch <- var_names$Z_names[var_switch$Z_names == 1]
  all_switch <- c(W1_switch, W2_switch, W3_switch, Z_switch)
    
  W1_names <- c(var_names$W1_names[var_names$W1_names == 0], 
    paste0(W1_switch, "0"), 
    paste0(W1_switch, "1"))
  W2_names <- c(var_names$W2_names[var_names$W2_names == 0], 
    paste0(W2_switch, "0"), 
    paste0(W2_switch, "1"))    
  W3_names <- c(var_names$W3_names[var_names$W3_names == 0], 
    paste0(W3_switch, "0"), 
    paste0(W3_switch, "1"))    
  Z_names <- c(var_names$Z_names[var_names$Z_names == 0], 
    paste0(Z_switch, "0"), 
    paste0(Z_switch, "1"))  
  
  # Adjust dataset accordingly:
  data[, paste0(all_switch, "0") := lapply(.SD, function(x)
    x * (1 - data[['switcher']])), .SDcols = all_switch]
  data[, paste0(all_switch, "1") := lapply(.SD, function(x)
    x * (1 - data[['switcher']])), .SDcols = all_switch]
}
